import os
import sys

from scipy.signal import savgol_filter

import torch
from orbitP.script import config
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from tqdm import tqdm
from datetime import datetime, timedelta

dis = ['R','S','W',"3D"]

def get_file_paths(directory):
    file_paths = []
    for root, dirs, files in os.walk(directory):
        for file in files:
            file_path = os.path.join(root, file)
            file_paths.append(os.path.normpath(file_path))
    file_paths = sorted(file_paths)
    return file_paths

def plot_loss(path_to_save, train=True):
    plt.rcParams.update({'font.size': 10})
    if train:
        with open(path_to_save+"save_loss/"+"train_loss.txt", 'r') as f:
            loss_list = [float(line.strip().split(": ")[1]) for line in f.readlines()]
    else:
        with open(path_to_save +"save_loss/"+"test_loss.txt", 'r') as f:
            loss_list = [float(line.strip().split(": ")[1]) for line in f.readlines()]
    if train:
        title = "Train"
    else:
        title = "Test"
    plt.plot(loss_list, label = "loss")
    plt.xlabel("Epochs")
    plt.ylabel("Loss")
    plt.legend()
    plt.title(title+"_loss")
    plt.savefig(path_to_save+f"save_predictions/{title}_loss.png")
    plt.close()

def plot_pml(path_to_save,pmlList,Title):
    plt.figure(figsize=(15, 6))
    plt.rcParams.update({"font.size": 16})
    x = [i * (1440/config.predicting_length) for i in range(1,len(pmlList)+1)]
    plt.plot(x, pmlList, '-', color='blue', label='Prediction Error', alpha=0.8)
    # 添加网格配置
    plt.grid(visible=True, which='major', linestyle='solid')  # 主网格线
    plt.minorticks_on()  # 打开次刻度
    plt.grid(visible=True, which='minor', linestyle='dashed', alpha=0.5)  # 次网格线
    plt.xlabel("t/min")
    plt.ylabel("Pml")
    plt.legend()
    plt.savefig(path_to_save+f"{Title}.png")
    plt.close()


def plot_origin_error(path_to_save, src, title):
    for idx,v in enumerate(dis):
        if idx>=config.outputSize:
            break
        if config.outputSize == 1:
            v = dis[config.outputIdx]
        plt.figure(figsize=(15,6))
        plt.rcParams.update({"font.size" : 16})
        # plotting
        x = [i * (1440/config.predicting_length) for i in range(1, len(src)+1)]
        plt.plot(x, src[:,idx], '-', color='red', label = 'origin', alpha=0.8)
        # 添加网格配置
        plt.grid(visible=True, which='major', linestyle='solid')  # 主网格线
        plt.minorticks_on()  # 打开次刻度
        plt.grid(visible=True, which='minor', linestyle='dashed', alpha=0.5)  # 次网格线
        plt.xlabel("t/min")
        plt.ylabel("error/m")
        plt.legend()
        plt.title(title+f" {v}")
        # save
        plt.savefig(path_to_save+f"{title} {v}.png")
        plt.close()

def plot_connetError(path_to_save, pre, suf, title):
    for idx,v in enumerate(dis):
        if idx>=config.outputSize:
            break
        if config.outputSize == 1:
            v = dis[config.outputIdx]
        plt.figure(figsize=(15,6))
        plt.rcParams.update({"font.size" : 16})
        # plotting
        x1 = [i * (1440/config.predicting_length) for i in range(1, len(pre)+1)]
        x2 = [i * (1440/config.predicting_length) for i in range(len(pre)+1, len(pre)+len(suf)+1)]
        plt.plot(x1, pre[:,idx], '-', color='red', label = 'pre', alpha=0.8)
        plt.plot(x2, suf[:,idx], '-', color='blue', label = 'suf', alpha=0.8)
        # 添加网格配置
        plt.grid(visible=True, which='major', linestyle='solid')  # 主网格线
        plt.minorticks_on()  # 打开次刻度
        plt.grid(visible=True, which='minor', linestyle='dashed', alpha=0.5)  # 次网格线
        plt.xlabel("t/min")
        plt.ylabel("error/m")
        plt.legend()
        plt.title(title+f" {v}")
        # save
        plt.savefig(path_to_save+f"{title} {v}.png")
        plt.close()

def plot_error(path_to_save, pred, src, title,length = config.predicting_length):
    for idx,v in enumerate(dis):
        if idx >= config.outputSize:
            break
        if config.outputSize == 1:
            v = dis[config.outputIdx]
        plt.figure(figsize=(15,6))
        plt.rcParams.update({"font.size" : 16})

        # x1 = [i * (1440/length) for i in range(1, len(src)//2+1)]
        # x2 = [i * (1440/length) for i in range(len(src)//2+1, len(src)+1)]
        # x3 = [i * (1440/length) for i in range(1, len(src)+1)]
        #
        # plt.plot(x1, src[:len(src)//2,idx], '-', color = 'red', label = 'origin preError', alpha=0.8)
        # plt.plot(x2, src[len(src)//2:,idx], '-', color = 'green', label = 'origin sufError', alpha=0.8)
        # plt.plot(x3, pred[:,idx], '-', color = 'blue', label = 'ewm smooth Error', alpha=0.8)


        # plotting
        x = [i * (1440/length) for i in range(1, len(src)+1)]
        plt.plot(x, src[:,idx], '-', color = 'red', label = 'True Error', alpha=0.8)
        plt.plot(x, pred[:,idx], '-', color = 'blue', label = 'Prediction Error', alpha=0.8)

        # 添加网格配置
        plt.grid(visible=True, which='major', linestyle='solid')  # 主网格线
        plt.minorticks_on()  # 打开次刻度
        plt.grid(visible=True, which='minor', linestyle='dashed', alpha=0.5)  # 次网格线
        plt.xlabel("t/min")
        plt.ylabel("error/m")
        # plt.ylim(bottom=0)
        plt.legend()
        plt.title(title+f" {v}")
        # save
        plt.savefig(path_to_save+f"{title} {v}.png")
        plt.close()

def plot_SULT(path_to_save, beginDate, endDate, title):
    dataPOEORB = {}
    filesPOEORBPath = get_file_paths(config.dataPOEORBDir)
    for filePOEORBPath in filesPOEORBPath:
        fileName = os.path.basename(filePOEORBPath).replace(".csv", "")
        fileDate = datetime.strptime(fileName, "%Y-%m-%d")
        if fileDate < beginDate or fileDate > (endDate + timedelta(days=1)):
            continue
        df_POEORB = pd.read_csv(filePOEORBPath)
        dataPOEORB_Time = df_POEORB.iloc[:, 0].to_numpy()
        dataPOEORB_chara = df_POEORB.iloc[:, 1:].to_numpy(dtype='float')
        for time, data in zip(dataPOEORB_Time, dataPOEORB_chara):
            timeDate = datetime.strptime(time, "%Y-%m-%d %H:%M:%S")
            dataPOEORB[timeDate] = data
    filesObserPath = get_file_paths(config.dataObsDir)
    t_bar = tqdm(filesObserPath, total=len(filesObserPath))
    for fileObserPath in t_bar:
        fileName = os.path.basename(fileObserPath).replace(".csv", "")
        fileDate = datetime.strptime(fileName, "%Y-%m-%d %H-%M-%S")
        t_bar.set_description(fileDate.strftime("%Y-%m-%d %H:%M:%S"))
        if fileDate < beginDate or fileDate > endDate:
            continue
        ObserNow = []
        MOEORBNow = []
        fileMOEORBPath = config.dataMOEORBDir + os.path.basename(fileObserPath)

        df_Obser = pd.read_csv(fileObserPath)
        dateObser = df_Obser['date'].to_numpy()
        df_MOEORB = pd.read_csv(fileMOEORBPath)
        dateMOEORB = df_MOEORB['date'].to_numpy()
        data_Obser = df_Obser[df_Obser.columns.drop('date')].to_numpy(dtype='float')
        data_MOEORB = df_MOEORB[df_MOEORB.columns.drop('date')].to_numpy(dtype='float')
        for time, data in zip(dateObser, data_Obser):
            timeDate = datetime.strptime(time, "%Y-%m-%d %H:%M:%S")
            # data[:config.outputSize] -= dataPOEORB[timeDate][:config.outputSize]
            err = data[:config.outputSize] - dataPOEORB[timeDate][:config.outputSize]
            data = np.hstack((err,data))
            ObserNow.append(data)
        for time, data in zip(dateMOEORB, data_MOEORB):
            timeDate = datetime.strptime(time, "%Y-%m-%d %H:%M:%S")
            # data[:config.outputSize] -= dataPOEORB[timeDate][:config.outputSize]
            err = data[:config.outputSize] - dataPOEORB[timeDate][:config.outputSize]
            data = np.hstack((err,data))
            MOEORBNow.append(data)
        ObserNow = np.array(ObserNow)
        MOEORBNow = np.array(MOEORBNow)

        for idx,v in enumerate(dis):
            if idx>=config.outputSize:
                break
            plt.figure(figsize=(15, 6))
            plt.rcParams.update({"font.size": 16})

            x1 = [i * (1440/config.predicting_length) for i in range(1, config.training_length+1)]
            x2 = [i * (1440/config.predicting_length) for i in range(config.training_length+1, 2*config.training_length+1)]
            plt.plot(x1, ObserNow[:,idx], '-', color = 'red', label = 'True Error', alpha=0.8)
            plt.scatter(x1, ObserNow[:,idx], color='black', s=2, label='points', alpha=0.9)
            plt.plot(x2, MOEORBNow[:,idx], '-', color = 'blue', label = 'Prediction Error', alpha=0.8)
            plt.scatter(x2, MOEORBNow[:,idx], color='black', s=2, label='points', alpha=0.9)

            # 添加网格配置
            plt.grid(visible=True, which='major', linestyle='solid')  # 主网格线
            plt.minorticks_on()  # 打开次刻度
            plt.grid(visible=True, which='minor', linestyle='dashed', alpha=0.5)  # 次网格线
            plt.xlabel("t/min")
            plt.ylabel("error/m")
            fileNewName = fileDate.strftime("%Y-%m-%d %H-%M-%S")
            plt.savefig(path_to_save+f"{title} {fileNewName} {v}.png")
            plt.close()

def plot_RSW(path_to_save, obsData, prdData, title, obsData_ori=None, prdData_ori=None):
    cnt = 0
    for idx,v in enumerate(dis):
        if idx>=config.outputSize:
            break
        obsList = range(0, len(obsData),config.training_length)
        prdList = range(0, len(prdData),config.predicting_length)
        t_bar = tqdm(obsList, total=len(obsList))
        for index,l in enumerate(t_bar):
            k = prdList[index]
            cnt+=1
            plt.figure(figsize=(15, 6))
            plt.rcParams.update({"font.size": 16})
            x1 = [i * (1440/config.predicting_length) for i in range(1, config.training_length+1)]
            x2 = [i * (1440/config.predicting_length) for i in range(config.training_length+1, config.training_length + config.predicting_length +1)]
            plt.plot(x1, obsData[l:l+config.training_length,idx], '-', color = 'red', label = 'True Error', alpha=0.8)
            plt.scatter(x1, obsData[l:l+config.training_length,idx], color='black', s=2, label='points', alpha=0.9)
            plt.plot(x2, prdData[k:k+config.predicting_length,idx], '-', color = 'blue', label = 'Prediction Error', alpha=0.8)
            plt.scatter(x2, prdData[k:k+config.predicting_length,idx], color='black', s=2, label='points', alpha=0.9)

            if obsData_ori is not None and prdData_ori is not None:
                plt.plot(x1, obsData_ori[l:l + config.training_length, idx], '-', color='orange', label='True Error',alpha=0.8)
                plt.scatter(x1, obsData_ori[l:l + config.training_length, idx], color='black', s=2, label='points',alpha=0.9)
                plt.plot(x2, prdData_ori[k:k+config.predicting_length,idx], '-', color='gray', label='Prediction Error',alpha=0.8)
                plt.scatter(x2, prdData_ori[k:k+config.predicting_length,idx], color='black', s=2, label='points',alpha=0.9)

            # 添加网格配置
            plt.grid(visible=True, which='major', linestyle='solid')  # 主网格线
            plt.minorticks_on()  # 打开次刻度
            plt.grid(visible=True, which='minor', linestyle='dashed', alpha=0.5)  # 次网格线
            plt.xlabel("t/min")
            plt.ylabel("error/m")
            # plt.legend()
            # plt.title(title+f" {v}")
            # save
            plt.savefig(path_to_save+f"{title}{l} {v} {cnt}.png")
            plt.close()
            if cnt>20:
                break

def plot_rmse(path_to_save, pred, src, title):
    for idx,v in enumerate(dis):
        if idx>=config.outputSize:
            break
        if config.outputSize == 1:
            v = dis[config.outputIdx]
        plt.figure(figsize=(15,6))
        plt.rcParams.update({"font.size" : 16})
        # plotting
        x = [i * (1440/config.predicting_length) for i in range(1, len(src)+1)]
        plt.plot(x, src[:,idx], '-', color = 'red', label = 'True Error', alpha=0.8)
        plt.plot(x, pred[:,idx], '-', color = 'blue', label = 'Prediction Error', alpha=0.8)

        # 添加网格配置
        plt.grid(visible=True, which='major', linestyle='solid')  # 主网格线
        plt.minorticks_on()  # 打开次刻度
        plt.grid(visible=True, which='minor', linestyle='dashed', alpha=0.5)  # 次网格线
        plt.xlabel("t/min")
        plt.ylabel("error/m")
        plt.ylim(bottom=0)
        plt.legend()
        plt.title(title+f" {v}")
        # save
        plt.savefig(path_to_save+f"{title} {v}.png")
        plt.close()

def plot_rmse_all(path_to_save, pred, title, ax):
    colors = ['red','blue','orange','green','purple','brown','pink','gray']
    plt.figure(figsize=(15, 6))
    plt.rcParams.update({"font.size": 16})
    # plotting
    x = [i * (1440 / config.predicting_length) for i in range(len(pred[0]))]
    for idx, predNow in enumerate(pred):
        predNow*=100
        plt.plot(x, predNow, '-', color=colors[idx], label=f'{title[idx]}', alpha=0.8)
        plt.scatter(x, predNow, color='black', s=2, alpha=0.5)

    # 添加网格配置
    plt.grid(visible=True, which='major', linestyle='solid')  # 主网格线
    plt.minorticks_on()  # 打开次刻度
    plt.grid(visible=True, which='minor', linestyle='dashed', alpha=0.5)  # 次网格线
    plt.xlabel("t/min")
    plt.ylabel("error/cm")
    plt.ylim(bottom=0)
    plt.legend()
    # save
    plt.savefig(path_to_save+f"all RMSE {dis[ax]}.png")
    plt.close()

def plot_mutilEval(path_to_save, meList, pmlList):
    meList = np.array(meList)
    pmlList = np.array(pmlList)
    num_mutilEval = meList.shape[0]
    colors = ['blue', 'orange', 'green', 'purple', 'brown', 'pink', 'gray']
    mutilEval_label = ['RMSE', 'MAE']
    num_label = len(mutilEval_label)
    x = np.arange(num_mutilEval)  # 类别编号 [0,1,2,...]
    bar_width = 0.25
    plt.figure(figsize=(25, 10))
    for i in range(num_label):
        plt.bar(x + i * bar_width, meList[:, i], width=bar_width, label=mutilEval_label[i],color=colors[i])
    plt.xticks(x + bar_width, [f'GPS {i+1}' for i in range(num_mutilEval)], fontsize=12)
    plt.ylabel("Metric Value", fontsize=12)
    plt.legend()
    plt.grid(axis='y', linestyle='--', alpha=0.6)
    plt.tight_layout()
    plt.savefig(path_to_save + f"mutilEval_RMSE.png")
    plt.close()

    mutilEval_label = ['Pml']
    num_label = len(mutilEval_label)
    x = np.arange(num_mutilEval)  # 类别编号 [0,1,2,...]
    bar_width = 0.25
    plt.figure(figsize=(25, 10))
    for i in range(num_label):
        plt.bar(x + i * bar_width, pmlList[i], width=bar_width, label=mutilEval_label[i],color=colors[i])
    plt.xticks(x + bar_width, [f'GPS {i+1}' for i in range(num_mutilEval)], fontsize=12)
    plt.ylabel("Metric Value", fontsize=12)
    plt.legend()
    plt.grid(axis='y', linestyle='--', alpha=0.6)
    plt.tight_layout()
    plt.savefig(path_to_save + f"mutilEval_Pml.png")
    plt.close()

def plot_rmse_3D(path_to_save, res, title):
    plt.figure(figsize=(15, 6))
    plt.rcParams.update({"font.size": 16})
    length = config.predicting_length
    # plotting
    x = [i * (1440 / length) for i in range(1, len(res) + 1)]
    rmse_3d = np.linalg.norm(res, axis=1)
    plt.plot(x, rmse_3d, '-', color='blue', label='residue', alpha=0.8)

    # 添加网格配置
    plt.grid(visible=True, which='major', linestyle='solid')  # 主网格线
    plt.minorticks_on()  # 打开次刻度
    plt.grid(visible=True, which='minor', linestyle='dashed', alpha=0.5)  # 次网格线
    plt.xlabel("t/min")
    plt.ylabel("error/m")
    plt.legend()
    # save
    plt.savefig(path_to_save + f"{title}.png")
    plt.close()

def plot_residue(path_to_save, pred, src, title):
    for idx,v in enumerate(dis):
        if idx>=config.outputSize:
            break
        if config.outputSize == 1:
            v = dis[config.outputIdx]
        plt.figure(figsize=(15,6))
        plt.rcParams.update({"font.size" : 16})

        # plotting
        x = [i * (1440/config.predicting_length) for i in range(1, len(src)+1)]
        residue  = [float(x) - float(y) for x, y in zip(pred[:,idx], src[:,idx])]
        plt.plot(x, residue, '-', color = 'blue', label = 'residue', alpha=0.8)

        # 添加网格配置
        plt.grid(visible=True, which='major', linestyle='solid')  # 主网格线
        plt.minorticks_on()  # 打开次刻度
        plt.grid(visible=True, which='minor', linestyle='dashed', alpha=0.5)  # 次网格线
        plt.xlabel("t/min")
        plt.ylabel("error/m")
        plt.legend()
        plt.title(title + f" {v}")
        # save
        plt.savefig(path_to_save+f"{title} {v}.png")
        plt.close()



def plot_training(epoch, path_to_save, src, prediction, sensor_number, index_in, index_tar):

    # idx_scr = index_in.tolist()[0]
    # idx_tar = index_tar.tolist()[0]
    # idx_pred = idx_scr.append(idx_tar.append([idx_tar[-1] + 1]))

    idx_scr = [i for i in range(len(src))]
    idx_pred = [i for i in range(1, len(prediction)+1)]

    plt.figure(figsize=(15,6))
    plt.rcParams.update({"font.size" : 18})
    plt.grid(b=True, which='major', linestyle = '-')
    plt.grid(b=True, which='minor', linestyle = '--', alpha=0.5)
    plt.minorticks_on()

    plt.plot(idx_scr, src, 'o-.', color = 'blue', label = 'input sequence', linewidth=1)
    plt.plot(idx_pred, prediction, 'o-.', color = 'limegreen', label = 'prediction sequence', linewidth=1)

    plt.title("Teaching Forcing from Sensor " + str(sensor_number[0]) + ", Epoch " + str(epoch))
    plt.xlabel("Time Elapsed")
    plt.ylabel("Humidity (%)")
    plt.legend()
    plt.savefig(path_to_save+f"/Epoch_{str(epoch)}.png")
    plt.close()

def plot_training_3(epoch, path_to_save, src, sampled_src, prediction, sensor_number, index_in, index_tar):

    # idx_scr = index_in.tolist()[0]
    # idx_tar = index_tar.tolist()[0]
    # idx_pred = idx_scr.append(idx_tar.append([idx_tar[-1] + 1]))

    idx_scr = [i for i in range(len(src))]
    idx_pred = [i for i in range(1, len(prediction)+1)]
    idx_sampled_src = [i for i in range(len(sampled_src))]

    plt.figure(figsize=(15,6))
    plt.rcParams.update({"font.size" : 18})
    plt.grid(b=True, which='major', linestyle = '-')
    plt.grid(b=True, which='minor', linestyle = '--', alpha=0.5)
    plt.minorticks_on()

    ## REMOVE DROPOUT FOR THIS PLOT TO APPEAR AS EXPECTED !! DROPOUT INTERFERES WITH HOW THE SAMPLED SOURCES ARE PLOTTED
    plt.plot(idx_sampled_src, sampled_src, 'o-.', color='red', label = 'sampled source', linewidth=1, markersize=10)
    plt.plot(idx_scr, src, 'o-.', color = 'blue', label = 'input sequence', linewidth=1)
    plt.plot(idx_pred, prediction, 'o-.', color = 'limegreen', label = 'prediction sequence', linewidth=1)
    plt.title("Teaching Forcing from Sensor " + str(sensor_number[0]) + ", Epoch " + str(epoch))
    plt.xlabel("Time Elapsed")
    plt.ylabel("Humidity (%)")
    plt.legend()
    plt.savefig(path_to_save+f"/Epoch_{str(epoch)}.png")
    plt.close()